package small.common.loadBalance.consisten;

import small.common.crypto.util.MD5Util;
import small.common.loadBalance.AbstractLoadBalance;
import small.common.loadBalance.BalanceInvoker;
import small.common.loadBalance.BalanceNode;

import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * 一致性hash
 *
 * @author jinzhou
 * @data 2023/4/2 15:44
 */
public class ConsistentHashLoadBalance<T extends BalanceNode, V extends BalanceInvoker> extends AbstractLoadBalance<T, V> {

    public final int REPLICA_NUMBER;

    public ConsistentHashLoadBalance() {
        this(160);
    }

    public ConsistentHashLoadBalance(int replicaNumber) {
        this.REPLICA_NUMBER = replicaNumber;
    }

    private final ConcurrentMap<String, ConsistentHashSelector<?>> selectors = new ConcurrentHashMap<String, ConsistentHashSelector<?>>();

    @Override
    protected T doSelect(List<T> nodes, V invoker) {

        String key = invoker.getKey();
        String consistentKey = invoker.getConsistentKey();
        int invokersHashCode = nodes.hashCode();
        ConsistentHashSelector<T> selector = (ConsistentHashSelector<T>) selectors.get(key);
        if (selector == null || selector.identityHashCode != invokersHashCode) {
            selectors.put(key, new ConsistentHashSelector<T>(nodes, invokersHashCode, REPLICA_NUMBER));
            selector = (ConsistentHashSelector<T>) selectors.get(key);
        }
        return selector.select(consistentKey);
    }


    private static final class ConsistentHashSelector<T extends BalanceNode> {
        /**
         * 虚拟节点
         */
        private final TreeMap<Long, T> virtualNodes;
        /**
         * 虚拟节点数
         */
        private final int replicaNumber;
        /**
         * 用于判断服务是否发生变化
         */
        private final int identityHashCode;


        ConsistentHashSelector(List<T> nodes, int identityHashCode, int replicaNumber) {
            this.virtualNodes = new TreeMap<Long, T>();
            this.identityHashCode = identityHashCode;
            this.replicaNumber = replicaNumber;
            for (T node : nodes) {
                String identifyId =  node.getIdentifyId();
                for (int i = 0; i < this.replicaNumber / 4; i++) {
                    // 对 identifyId + i 进行 md5 运算，得到一个长度为16的字节数组
                    byte[] digest = MD5Util.getMD5Bytes(identifyId + i);
                    // 对 digest 部分字节进行4次 hash 运算，得到四个不同的 long 型正整数
                    for (int h = 0; h < 4; h++) {
                        // h = 0 时，取 digest 中下标为 0 ~ 3 的4个字节进行位运算
                        // h = 1 时，取 digest 中下标为 4 ~ 7 的4个字节进行位运算
                        // h = 2, h = 3 时过程同上
                        long m = hash(digest, h);
                        // 将 hash 到 invoker 的映射关系存储到 virtualInvokers 中，
                        // virtualNodes 需要提供高效的查询操作，因此选用 TreeMap 作为存储结构
                        virtualNodes.put(m, node);
                    }
                }
            }
        }

        public T select(String consistentKey) {
            byte[] digest = MD5Util.getMD5Bytes(consistentKey);
            // 取 digest 数组的前四个字节进行 hash 运算，再将 hash 值传给 selectForKey 方法，
            // 寻找合适的 node
            return selectForKey(hash(digest, 0));
        }


        private T selectForKey(long hash) {
            Map.Entry<Long, T> entry = virtualNodes.ceilingEntry(hash);
            if (entry == null) {
                entry = virtualNodes.firstEntry();
            }
            return entry.getValue();
        }

        private long hash(byte[] digest, int number) {
            return (((long) (digest[3 + number * 4] & 0xFF) << 24)
                    | ((long) (digest[2 + number * 4] & 0xFF) << 16)
                    | ((long) (digest[1 + number * 4] & 0xFF) << 8)
                    | (digest[number * 4] & 0xFF))
                    & 0xFFFFFFFFL;
        }
    }
}
